import os
import shutil
import sys
from pathlib import Path

def process_model_folders():
    """处理每个型号文件夹，仅保留一个DAT文件（本地环境版）"""
    try:
        # 定义源目录和目标目录
        source_dir = "data_2数据清洗（品牌,型号）"
        target_dir = "data_2数据清洗（品牌,型号，去除重复）"
        
        # 验证目录是否存在
        if not os.path.exists(source_dir):
            print(f"Error: 源目录不存在: {source_dir}", file=sys.stderr)
            return False
            
        # 创建目标目录（如果不存在）
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
        
        # 统计信息
        stats = {
            "total_brands_processed": 0,
            "total_models_processed": 0,
            "total_dat_files_kept": 0,
            "models_without_dat": []
        }
        
        # 识别有效品牌（包含型号文件夹的品牌）
        valid_brands = []
        for brand_name in os.listdir(source_dir):
            brand_path = os.path.join(source_dir, brand_name)
            if not os.path.isdir(brand_path):
                continue
                
            # 检查是否有型号文件夹
            model_folders = []
            for item in os.listdir(brand_path):
                item_path = os.path.join(brand_path, item)
                if os.path.isdir(item_path):
                    model_folders.append(item)
                    
            if model_folders:
                valid_brands.append({
                    'brand_name': brand_name,
                    'brand_path': brand_path,
                    'model_folders': model_folders
                })
        
        # 遍历每个有效品牌
        for brand in valid_brands:
            brand_name = brand['brand_name']
            brand_path = brand['brand_path']
            model_folders = brand['model_folders']
            
            stats["total_brands_processed"] += 1
            
            # 遍历品牌下的每个型号文件夹
            for model_name in model_folders:
                model_source_path = os.path.join(brand_path, model_name)
                model_target_path = os.path.join(target_dir, brand_name, model_name)
                
                # 创建目标型号文件夹
                os.makedirs(model_target_path, exist_ok=True)
                
                # 查找所有DAT文件（不区分大小写）
                dat_files = []
                for item in os.listdir(model_source_path):
                    item_path = os.path.join(model_source_path, item)
                    if os.path.isfile(item_path) and item.lower().endswith('.dat'):
                        dat_files.append(item)
                
                # 按文件名排序，确保选择一致性
                dat_files.sort()
                
                # 处理DAT文件
                if dat_files:
                    # 只保留第一个DAT文件
                    selected_dat = dat_files[0]
                    source_dat_path = os.path.join(model_source_path, selected_dat)
                    target_dat_path = os.path.join(model_target_path, selected_dat)
                    
                    # 复制文件（如果已存在则跳过）
                    if not os.path.exists(target_dat_path):
                        shutil.copy2(source_dat_path, target_dat_path)
                    stats["total_dat_files_kept"] += 1
                else:
                    # 记录没有DAT文件的型号
                    stats["models_without_dat"].append(f"{brand_name}/{model_name}")
                
                stats["total_models_processed"] += 1
        
        # 输出结果消息
        result_message = (
            f"型号文件夹处理完成！\n"
            f"处理品牌数: {stats['total_brands_processed']}\n"
            f"处理型号数: {stats['total_models_processed']}\n"
            f"保留DAT文件数: {stats['total_dat_files_kept']}\n"
        )
        
        if stats["models_without_dat"]:
            result_message += (
                f"注意：{len(stats['models_without_dat'])}个型号文件夹没有找到DAT文件:\n"
                f"{', '.join(stats['models_without_dat'][:5])}{'...' if len(stats['models_without_dat'])>5 else ''}"
            )
        
        print(result_message)
        return True
        
    except Exception as e:
        error_msg = f"处理型号文件夹时出错: {str(e)}"
        print(error_msg, file=sys.stderr)
        return False

# 执行处理
if __name__ == "__main__":
    process_model_folders()
